#include "util/graph_io.hpp"

//#define TEST

#include "baseline.hpp"
#include "util/vecs_io.hpp"
#include "util/ground_truth.hpp"
#include "util/time_memory.hpp"

using namespace std;
using namespace GraphGradient;


int main() {
    const char *dataset_name = "sift";
    const char *graph_method = "hnsw";
    const int M = 40;
    const int ef_construction = 100;
    const int random_seed = 100;

    int ef_search_config[3] = {10, 2, 20}; //start, stride, iter
    /*
     * gist
     * int ef_search_config[3] = {50, 3, 60}; //start, stride, iter
     */

    int n_query, top_k, n_base, data_dim;
    top_k = 10;

    vector<float *> data = read_data(dataset_name, n_query, top_k, n_base, data_dim);
    float *base = data[0];
    float *query = data[1];
    int *gnd = (int *) data[2];

    printf("dataset %s, M %d, efConstruction %d\n", dataset_name, M, ef_construction);

    //build index
    printf("building index\n");
    int indexing_time;
    TimeRecord total_insert_record;
    GraphGradient::GraphBaseline graphBaseline(data_dim, n_base, random_seed + 1);
    {
        char path[256];
        sprintf(path, "../data/%s/%s.graph", dataset_name, graph_method);
        vector <vector<int>> graph;
        int enter_point;
        read_graph<int>(path, graph, enter_point);

        graphBaseline.add(base, graph, enter_point);
        for (int i = 0; i < n_base; i++) {
            graph[i].clear();
        }
        graph.clear();
    }
    indexing_time = (int) (total_insert_record.get_elapsed_time_micro() * 1e-6);
    total_insert_record.reset();

    printf("warming up\n");
    for (int iter = 0, ef_search = ef_search_config[0];
         iter < 12; iter++, ef_search += 1) {
        TimeRecord query_record;
        vector <unordered_set<int>> test_gnd_l = graphBaseline.search_knn(query, top_k, n_query, ef_search);
        double query_time = query_record.get_elapsed_time_micro() / n_query * 1e-3; //unit ms

        vector<float> tmp_part_recall_l(n_query);
        double recall = count_recall(n_query, top_k, test_gnd_l, gnd, tmp_part_recall_l);
        printf("efSearch %d, average recall: %.3f, query time %.3f\n", ef_search, recall, query_time);
    }
#ifdef TEST
    graphBaseline.is_warm_up_ = false;
#endif
    vector <pair<double, double>> time_recall_l(ef_search_config[2]);
    vector<int> ef_search_l(ef_search_config[2]);
    vector <vector<float>> recall_part_l_l(ef_search_config[2]);

    printf("querying\n");
    for (int iter = 0, ef_search = ef_search_config[0];
         iter < ef_search_config[2]; iter++, ef_search += ef_search_config[1]) {
        ef_search_l[iter] = ef_search;

        TimeRecord query_record;
        vector <unordered_set<int>> test_gnd_l = graphBaseline.search_knn(query, top_k, n_query, ef_search);
        double query_time = query_record.get_elapsed_time_micro() / n_query * 1e-3; //unit ms

        recall_part_l_l[iter] = vector<float>(n_query);
        double recall = count_recall(n_query, top_k, test_gnd_l, gnd, recall_part_l_l[iter]);
#ifndef TEST
        printf("efSearch %d, average recall: %.3f, query time %.3f ms\n", ef_search, recall, query_time);
#else
        printf("efSearch %d, average recall: %.3f, query time %.3f ms, n_hops %d, n_visit_edges %d, n_visited_edges %d\n",
               ef_search,
               recall, query_time, graphBaseline.test_avg_n_hops_l.back(),
               graphBaseline.test_avg_n_visit_edges_l.back(), graphBaseline.test_avg_n_visited_edges_l.back());

#endif
        time_recall_l[iter] = make_pair(query_time, recall);
    }

    size_t peak_memory = get_peak_RSS();
    printf("peak memory %ld Mb\n", peak_memory / 1000000);

    SaveNode save_node;
    save_node.dataset = dataset_name;
    save_node.method_name = "baseline";
    save_node.graph_method = graph_method;
    save_node.top_k = top_k;
    save_node.time_recall_l = time_recall_l;
    save_node.recall_part_l_l = recall_part_l_l;
    save_node.ef_search_l = ef_search_l;
    save_node.n_index = 1;
    save_node.total_M = M;
    save_node.ef_construction = ef_construction;
    save_node.indexing_time_l = {indexing_time};
    save_node.max_memory = peak_memory / 1000000;
    save_node.M_l = {M};
#ifdef TEST
    string avg_n_hops = save_node.array_to_string(graphBaseline.test_avg_n_hops_l);
    save_node.other_config.insert(pair<string, string>("n_hops", avg_n_hops));

    string avg_n_visit_edges = save_node.array_to_string(graphBaseline.test_avg_n_visit_edges_l);
    save_node.other_config.insert(pair<string, string>("n_visit_edges", avg_n_visit_edges));

    string avg_n_visited_edges = save_node.array_to_string(graphBaseline.test_avg_n_visited_edges_l);
    save_node.other_config.insert(pair<string, string>("n_visited_edges", avg_n_visited_edges));
#endif

    save_node.save();

    return 0;
}